"""
RAG管理器
负责检索增强生成功能
"""

from typing import List, Dict, Any, Optional
import chromadb
from chromadb.config import Settings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
import numpy as np
import hashlib

from app.models.knowledge_base import KnowledgeBase
from app.core.config import settings


class RAGManager:
    """RAG管理器"""
    
    def __init__(self):
        self.chroma_client = chromadb.PersistentClient(
            path=settings.chroma_persist_directory
        )
        # 使用简单的文本哈希作为嵌入的替代方案
        self.embedding_dimension = 1536
    
    def _generate_simple_embedding(self, text: str) -> List[float]:
        """生成简单的文本嵌入向量"""
        # 使用文本哈希生成确定性的向量
        text_hash = hashlib.md5(text.encode('utf-8')).hexdigest()
        
        # 将哈希转换为数值
        hash_int = int(text_hash[:8], 16)
        
        # 生成确定性的向量
        np.random.seed(hash_int)
        embedding = np.random.rand(self.embedding_dimension).tolist()
        
        return embedding
    
    def create_collection(self, knowledge_base: KnowledgeBase) -> bool:
        """创建向量集合"""
        try:
            # 检查集合是否已存在
            existing_collections = [col.name for col in self.chroma_client.list_collections()]
            if knowledge_base.collection_name in existing_collections:
                return True
            
            # 创建新集合
            self.chroma_client.create_collection(
                name=knowledge_base.collection_name,
                metadata={"description": knowledge_base.description}
            )
            return True
        except Exception as e:
            return False
    
    def add_documents(
        self, 
        knowledge_base: KnowledgeBase, 
        documents: List[Document]
    ) -> bool:
        """添加文档到向量数据库"""
        try:
            # 确保集合存在
            if not self.create_collection(knowledge_base):
                return False
            
            # 获取集合
            collection = self.chroma_client.get_collection(knowledge_base.collection_name)
            
            # 准备数据
            texts = [doc.page_content for doc in documents]
            metadatas = [doc.metadata for doc in documents]
            ids = [f"doc_{i}_{hash(doc.page_content) % 100000}" for i, doc in enumerate(documents)]
            
            # 生成嵌入向量
            embeddings = [self._generate_simple_embedding(text) for text in texts]
            
            # 添加到集合
            collection.add(
                embeddings=embeddings,
                documents=texts,
                metadatas=metadatas,
                ids=ids
            )
            
            return True
        except Exception as e:
            return False
    
    async def retrieve_documents(
        self, 
        knowledge_base: KnowledgeBase, 
        query: str, 
        top_k: int = 5
    ) -> List[Document]:
        """检索相关文档"""
        try:
            # 获取集合
            collection = self.chroma_client.get_collection(knowledge_base.collection_name)
            
            # 生成查询嵌入
            query_embedding = self._generate_simple_embedding(query)
            
            # 检索相似文档
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=top_k
            )
            
            # 构建Document对象
            documents = []
            if results['documents'] and results['documents'][0]:
                for i, content in enumerate(results['documents'][0]):
                    metadata = results['metadatas'][0][i] if results['metadatas'] and results['metadatas'][0] else {}
                    metadata['distance'] = results['distances'][0][i] if results['distances'] and results['distances'][0] else 0
                    
                    documents.append(Document(
                        page_content=content,
                        metadata=metadata
                    ))
            
            return documents
        except Exception as e:
            return []
    
    def delete_collection(self, knowledge_base: KnowledgeBase) -> bool:
        """删除向量集合"""
        try:
            self.chroma_client.delete_collection(knowledge_base.collection_name)
            return True
        except Exception as e:
            return False
    
    def get_collection_info(self, knowledge_base: KnowledgeBase) -> Dict[str, Any]:
        """获取集合信息"""
        try:
            collection = self.chroma_client.get_collection(knowledge_base.collection_name)
            count = collection.count()
            
            return {
                "name": knowledge_base.collection_name,
                "document_count": count,
                "exists": True
            }
        except Exception:
            return {
                "name": knowledge_base.collection_name,
                "document_count": 0,
                "exists": False
            }
    
    def split_documents(
        self, 
        documents: List[Document], 
        chunk_size: int = 1000, 
        chunk_overlap: int = 200
    ) -> List[Document]:
        """分割文档"""
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            separators=["\n\n", "\n", " ", ""]
        )
        
        return text_splitter.split_documents(documents)
    
    def build_retrieval_chain(
        self, 
        knowledge_base: KnowledgeBase, 
        llm
    ):
        """构建检索链"""
        try:
            # 直接使用Chroma客户端进行检索
            collection = self.chroma_client.get_collection(knowledge_base.collection_name)
            
            def retriever_func(query: str, k: int = 5):
                """检索函数"""
                query_embedding = self._generate_simple_embedding(query)
                results = collection.query(
                    query_embeddings=[query_embedding],
                    n_results=k
                )
                
                documents = []
                if results['documents'] and results['documents'][0]:
                    for i, content in enumerate(results['documents'][0]):
                        metadata = results['metadatas'][0][i] if results['metadatas'] and results['metadatas'][0] else {}
                        documents.append(Document(
                            page_content=content,
                            metadata=metadata
                        ))
                
                return documents
            
            return retriever_func
        except Exception as e:
            return None
